{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention for seq2seq\n",
    "\n",
    "This notebook has been adapted fom the pytorch tutorial [NLP FROM SCRATCH: TRANSLATION WITH A SEQUENCE TO SEQUENCE NETWORK AND ATTENTION](https://pytorch.org/tutorials/intermediate/seq2seq_translation_tutorial.html) by Sean Robertson and adapted by Marc Lelarge for the [deep learning course](https://dataflowr.github.io/website/)\n",
    "\n",
    "**(02/2023) There are important differences with the tutorial.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "%matplotlib inline\n",
    "from __future__ import unicode_literals, print_function, division\n",
    "from io import open\n",
    "import unicodedata\n",
    "import string\n",
    "import re\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch import optim\n",
    "import torch.nn.functional as F\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# where data is stored:\n",
    "# if on COLAB, you can comment the following lines \n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "ROOT_DIR = Path.home()\n",
    "data_path = os.path.join(ROOT_DIR,'data/')\n",
    "\n",
    "# and uncomment the following lines if on COLAB, you can download the data with:\n",
    "#!wget https://download.pytorch.org/tutorial/data.zip\n",
    "#!unzip data.zip\n",
    "#data_path = './'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def running_mean(x, N=100):\n",
    "    cumsum = np.cumsum(np.insert(x, 0, 0)) \n",
    "    return (cumsum[N:] - cumsum[:-N]) / float(N)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# data preprocessing\n",
    "\n",
    "This code is directly taken from the Pytorch tutorial and creates the corpus of pairs of sentences in french and english as well as the french and englsih tokenizers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SOS_token = 0\n",
    "EOS_token = 1\n",
    "\n",
    "class Lang:\n",
    "    def __init__(self, name):\n",
    "        self.name = name\n",
    "        self.word2index = {}\n",
    "        self.word2count = {}\n",
    "        self.index2word = {0: \"SOS\", 1: \"EOS\"}\n",
    "        self.n_words = 2  # Count SOS and EOS\n",
    "\n",
    "    def addSentence(self, sentence):\n",
    "        for word in sentence.split(' '):\n",
    "            self.addWord(word)\n",
    "\n",
    "    def addWord(self, word):\n",
    "        if word not in self.word2index:\n",
    "            self.word2index[word] = self.n_words\n",
    "            self.word2count[word] = 1\n",
    "            self.index2word[self.n_words] = word\n",
    "            self.n_words += 1\n",
    "        else:\n",
    "            self.word2count[word] += 1\n",
    "\n",
    "# Turn a Unicode string to plain ASCII, thanks to\n",
    "# https://stackoverflow.com/a/518232/2809427\n",
    "def unicodeToAscii(s):\n",
    "    return ''.join(\n",
    "        c for c in unicodedata.normalize('NFD', s)\n",
    "        if unicodedata.category(c) != 'Mn'\n",
    "    )\n",
    "\n",
    "# Lowercase, trim, and remove non-letter characters\n",
    "\n",
    "def normalizeString(s):\n",
    "    s = unicodeToAscii(s.lower().strip())\n",
    "    s = re.sub(r\"([.!?])\", r\" \\1\", s)\n",
    "    s = re.sub(r\"[^a-zA-Z.!?]+\", r\" \", s)\n",
    "    return s\n",
    "\n",
    "def readLangs(lang1, lang2, reverse=False):\n",
    "    print(\"Reading lines...\")\n",
    "\n",
    "    # Read the file and split into lines\n",
    "    lines = open(data_path+'data/%s-%s.txt' % (lang1, lang2), encoding='utf-8').\\\n",
    "        read().strip().split('\\n')\n",
    "\n",
    "    # Split every line into pairs and normalize\n",
    "    pairs = [[normalizeString(s) for s in l.split('\\t')] for l in lines]\n",
    "\n",
    "    # Reverse pairs, make Lang instances\n",
    "    if reverse:\n",
    "        pairs = [list(reversed(p)) for p in pairs]\n",
    "        input_lang = Lang(lang2)\n",
    "        output_lang = Lang(lang1)\n",
    "    else:\n",
    "        input_lang = Lang(lang1)\n",
    "        output_lang = Lang(lang2)\n",
    "\n",
    "    return input_lang, output_lang, pairs\n",
    "\n",
    "MAX_LENGTH = 10\n",
    "\n",
    "eng_prefixes = (\n",
    "    \"i am \", \"i m \",\n",
    "    \"he is\", \"he s \",\n",
    "    \"she is\", \"she s \",\n",
    "    \"you are\", \"you re \",\n",
    "    \"we are\", \"we re \",\n",
    "    \"they are\", \"they re \"\n",
    ")\n",
    "\n",
    "def filterPair(p):\n",
    "    return len(p[0].split(' ')) < MAX_LENGTH and \\\n",
    "        len(p[1].split(' ')) < MAX_LENGTH and \\\n",
    "        p[1].startswith(eng_prefixes)\n",
    "\n",
    "def filterPair_test(p):\n",
    "    return len(p[0].split(' ')) > MAX_LENGTH and \\\n",
    "        len(p[1].split(' ')) > MAX_LENGTH and \\\n",
    "        p[1].startswith(eng_prefixes)\n",
    "\n",
    "def filterPairs(pairs):\n",
    "    return [pair for pair in pairs if filterPair(pair)]\n",
    "\n",
    "def filterPairs_test(pairs):\n",
    "    return [pair for pair in pairs if filterPair_test(pair)]\n",
    "\n",
    "def prepareData(lang1, lang2, reverse=False):\n",
    "    input_lang, output_lang, pairs = readLangs(lang1, lang2, reverse)\n",
    "    print(\"Read %s sentence pairs\" % len(pairs))\n",
    "    pairs = filterPairs(pairs)\n",
    "    print(\"Trimmed to %s sentence pairs\" % len(pairs))\n",
    "    print(\"Counting words...\")\n",
    "    for pair in pairs:\n",
    "        input_lang.addSentence(pair[0])\n",
    "        output_lang.addSentence(pair[1])\n",
    "    print(\"Counted words:\")\n",
    "    print(input_lang.name, input_lang.n_words)\n",
    "    print(output_lang.name, output_lang.n_words)\n",
    "    return input_lang, output_lang, pairs\n",
    "\n",
    "input_lang, output_lang, pairs = prepareData('eng', 'fra', True)\n",
    "print(random.choice(pairs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pairs_train, pairs_val = train_test_split(pairs, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def indexesFromSentence(lang, sentence):\n",
    "    return [lang.word2index[word] for word in sentence.split(' ')]\n",
    "\n",
    "def tensorFromSentence(lang, sentence):\n",
    "    indexes = indexesFromSentence(lang, sentence)\n",
    "    indexes.append(EOS_token)\n",
    "    return torch.tensor(indexes, dtype=torch.long, device=device).view(-1, 1)\n",
    "\n",
    "def tensorsFromPair(pair):\n",
    "    input_tensor = tensorFromSentence(input_lang, pair[0])\n",
    "    target_tensor = tensorFromSentence(output_lang, pair[1])\n",
    "    return (input_tensor, target_tensor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_pairs = [tensorsFromPair(pv) for pv in pairs_val]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_pairs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seq2seq\n",
    "\n",
    "Here we follow the PyTorch tutorial and implement [Sequence to Sequence Learning with Neural Networks](https://arxiv.org/abs/1409.3215v3). The only modification in the code is that the encoder takes a full sentence (and no hidden state) in the forward pass and outputs all the correponding hidden states. As a result, there is no need to make a for loop for the encoder. Still, to get things simple, we do not deal with batches, if you want to deal with batches have a look at [Batches with sequences in Pytorch](https://dataflowr.github.io/website/modules/11c-batches-with-sequences/).\n",
    "\n",
    "We also train on a test set and compute the loss on a validation set (see the split done on the corpus above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderRNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size):\n",
    "        super(EncoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = nn.Embedding(input_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size)\n",
    "\n",
    "    def forward(self, input):\n",
    "        embedded = self.embedding(input)\n",
    "        output, _ = self.gru(embedded, self.initHidden())\n",
    "        return output\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_iters = 4\n",
    "training_pairs = [tensorsFromPair(random.choice(pairs))\n",
    "                      for i in range(n_iters)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_pairs[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_pairs[0][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "one_input = training_pairs[0][0]\n",
    "out = encoder(one_input)\n",
    "# here is the code used in the tutorial:\n",
    "#encoder_hidden = encoder.initHidden()\n",
    "#for c in one_input:\n",
    "#    out, encoder_hidden = encoder(c,encoder_hidden)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super(DecoderRNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.embedding = nn.Embedding(output_size, hidden_size)\n",
    "        self.gru = nn.GRU(hidden_size, hidden_size)\n",
    "        self.out = nn.Linear(hidden_size, output_size)\n",
    "        #self.softmax = nn.LogSoftmax(dim=1) find the notation misleading...\n",
    "\n",
    "    def forward(self, input, hidden, *kwargs):\n",
    "        output = self.embedding(input).view(1, 1, -1)\n",
    "        #output = F.relu(output) Why taking a relu on embeddings?\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        #output = self.softmax(self.out(output[0]))\n",
    "        return F.log_softmax(self.out(output[0]),dim=1), hidden, None #do not worry about last argument!\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = DecoderRNN(hidden_size, output_lang.n_words).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "decoder_hidden = out[-1,:,:].view(1,1,hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_hidden.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output, hidden, _ = decoder(decoder_input, decoder_hidden)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_onepair(input_tensor, target_tensor, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, teacher_forcing_ratio = 0.5, with_attention = False):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "    \n",
    "    encoder = encoder.train()\n",
    "    decoder = decoder.train()\n",
    "    encoder_optimizer.zero_grad()\n",
    "    decoder_optimizer.zero_grad()\n",
    "    target_length = target_tensor.size(0)\n",
    "    loss = 0\n",
    "    encoder_outputs = encoder(input_tensor)\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "    \n",
    "    if with_attention:\n",
    "        decoder_hidden = decoder.initHidden()\n",
    "    else:\n",
    "        decoder_hidden = encoder_outputs[-1,:,:].view(1,1,hidden_size)\n",
    "    \n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    loss.backward()\n",
    "    encoder_optimizer.step()\n",
    "    decoder_optimizer.step()\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_onepair(input_tensor, target_tensor, encoder, decoder, criterion, teacher_forcing_ratio = 0.5, with_attention = False):\n",
    "    encoder_hidden = encoder.initHidden()\n",
    "\n",
    "    encoder = encoder.eval()\n",
    "    decoder = decoder.eval()\n",
    "    target_length = target_tensor.size(0)\n",
    "    loss = 0\n",
    "    encoder_outputs = encoder(input_tensor)\n",
    "    decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "    \n",
    "    if with_attention:\n",
    "        decoder_hidden = decoder.initHidden()\n",
    "    else:\n",
    "        decoder_hidden = encoder_outputs[-1,:,:].view(1,1,hidden_size)\n",
    "        \n",
    "    use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False\n",
    "\n",
    "    if use_teacher_forcing:\n",
    "        # Teacher forcing: Feed the target as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            decoder_input = target_tensor[di]  # Teacher forcing\n",
    "\n",
    "    else:\n",
    "        # Without teacher forcing: use its own predictions as the next input\n",
    "        for di in range(target_length):\n",
    "            decoder_output, decoder_hidden, _ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            topv, topi = decoder_output.topk(1)\n",
    "            decoder_input = topi.squeeze().detach()  # detach from history as input\n",
    "            loss += criterion(decoder_output, target_tensor[di])\n",
    "            if decoder_input.item() == EOS_token:\n",
    "                break\n",
    "\n",
    "    return loss.item() / target_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainIters(encoder, decoder, n_iters, print_every=1000, learning_rate=0.01, teacher_forcing_ratio=0.9, with_attention = False):\n",
    "    \n",
    "    plot_losses = []\n",
    "    plot_losses_val = []\n",
    "    print_loss_total = 0  # Reset every print_every\n",
    "\n",
    "    encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "    decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "    training_pairs = [tensorsFromPair(random.choice(pairs_train))\n",
    "                      for i in range(n_iters)]\n",
    "    criterion = nn.NLLLoss()\n",
    "\n",
    "    for iter in range(1, n_iters + 1):\n",
    "        training_pair = training_pairs[iter - 1]\n",
    "        input_tensor = training_pair[0]\n",
    "        target_tensor = training_pair[1]\n",
    "\n",
    "        loss = train_onepair(input_tensor, target_tensor, encoder,\n",
    "                     decoder, encoder_optimizer, decoder_optimizer, criterion, teacher_forcing_ratio, with_attention)\n",
    "        print_loss_total += loss\n",
    "        plot_losses.append(loss)\n",
    "\n",
    "        if iter % print_every == 0:\n",
    "            loss_val = 0\n",
    "            for (input_tensor, target_tensor) in val_pairs:\n",
    "                loss = eval_onepair(input_tensor, target_tensor, encoder,\n",
    "                     decoder, criterion, teacher_forcing_ratio, with_attention)\n",
    "                loss_val += loss\n",
    "            loss_val = loss_val/len(val_pairs)\n",
    "            plot_losses_val.append(loss_val)\n",
    "            print_loss_avg = print_loss_total / print_every\n",
    "            print_loss_total = 0\n",
    "            print('(%d %d%%) loss train %.4f and val %.4f' % (iter, iter / n_iters * 100, print_loss_avg, loss_val))\n",
    "    return plot_losses, plot_losses_val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate=0.01\n",
    "encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "criterion = nn.NLLLoss()\n",
    "train_onepair(training_pairs[0][0],training_pairs[0][1],encoder,decoder,encoder_optimizer,decoder_optimizer,criterion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 128\n",
    "n_epochs = 50000\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "decoder = DecoderRNN(hidden_size, output_lang.n_words).to(device)\n",
    "plot_losses,plot_losses_val = trainIters(encoder,decoder,n_epochs,print_every=2500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(running_mean(plot_losses))\n",
    "plt.plot([2500*i for i in range(int(n_epochs/2500))], plot_losses_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(encoder, decoder, sentence, max_length=MAX_LENGTH, with_attention = False):\n",
    "    with torch.no_grad():\n",
    "        input_tensor = tensorFromSentence(input_lang, sentence)\n",
    "        input_length = input_tensor.size()[0]\n",
    "        encoder_hidden = encoder.initHidden()\n",
    "\n",
    "        encoder_outputs = encoder(input_tensor)\n",
    "\n",
    "        decoder_input = torch.tensor([[SOS_token]], device=device)  # SOS\n",
    "\n",
    "        if with_attention:\n",
    "            decoder_hidden = decoder.initHidden()\n",
    "        else:\n",
    "            decoder_hidden = encoder_outputs[-1,:,:].view(1,1,hidden_size)\n",
    "        \n",
    "        decoded_words = []\n",
    "\n",
    "        for di in range(max_length):\n",
    "            decoder_output, decoder_hidden,_ = decoder(decoder_input, decoder_hidden, encoder_outputs)\n",
    "            \n",
    "            topv, topi = decoder_output.data.topk(1)\n",
    "            if topi.item() == EOS_token:\n",
    "                decoded_words.append('<EOS>')\n",
    "                break\n",
    "            else:\n",
    "                decoded_words.append(output_lang.index2word[topi.item()])\n",
    "\n",
    "            decoder_input = topi.squeeze().detach()\n",
    "\n",
    "        return decoded_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluateRandomly(encoder, decoder, n=10, with_attention = False):\n",
    "    for i in range(n):\n",
    "        pair = random.choice(pairs_val)\n",
    "        print('>', pair[0])\n",
    "        print('=', pair[1])\n",
    "        output_words = evaluate(encoder, decoder, pair[0], with_attention =with_attention)\n",
    "        output_sentence = ' '.join(output_words)\n",
    "        print('<', output_sentence)\n",
    "        print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateRandomly(encoder, decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attention mechanism\n",
    "\n",
    "Here we implement the attention mechanism inspired from the paper [Neural Machine Translation by Jointly Learning to Align and Translate](https://arxiv.org/abs/1409.0473). \n",
    "\n",
    "![](attention_bahdanau.png)\n",
    "\n",
    "\n",
    "For the alignment model $a$, we take $e_{ij} = v_a^T\\text{tanh}(W_a[s_{i-1}, h_j])$. So to compute the attention weights $(\\alpha_{ij})_{j=1}^T$, we first build the $[s_{i-1}, h_1], \\dots, [s_{i-1}, h_T]$ by repeating the hidden state $s_{i-1}$ and concatenating it with the approriate $h_j$. We obtain a tensor of size $T\\times 1\\times 2H$ (because we're dealing with batches of size 1) that is used as input for the alignment model. After applying a softmax, we get the attention weights.\n",
    "\n",
    "The code below is siginficantly different from the original Pytorch Tutorial..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttnDecoderRNN(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size):\n",
    "        super().__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.embedding = nn.Embedding(self.output_size, self.hidden_size, padding_idx=0)\n",
    "        self.attn_w = nn.Linear(2 * self.hidden_size, self.hidden_size)\n",
    "        self.attn_v = nn.Linear(self.hidden_size, 1)\n",
    "        self.gru = nn.GRU(self.hidden_size * 2, self.hidden_size)\n",
    "        self.out = nn.Linear(self.hidden_size * 2, self.output_size)\n",
    "\n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "        # B = 1 batch size\n",
    "        # encoder_outputs (L,B,H)\n",
    "        seq_len, _, _ = encoder_outputs.shape\n",
    "        # hidden (1,B,H)\n",
    "        hidden_tile = hidden.repeat(seq_len, 1, 1) # (L,B,H)\n",
    "        # concat \n",
    "        concat = torch.cat((hidden_tile, encoder_outputs), dim=2) # (L,B,2*H)\n",
    "        attn_weights = F.softmax(self.attn_v(torch.tanh(self.attn_w(concat))), dim=0)\n",
    "        # attn_weights (L,B,1) torch.sum(attn_weights,dim=0) = 1\n",
    "        contexts = torch.einsum('jbi,jbk->ibk', attn_weights, encoder_outputs) # (1,B,H)\n",
    "        # input is previous prediction \n",
    "        embedded = self.embedding(input).view(1,1,-1) # (1,B,H)\n",
    "        output = torch.cat((embedded, contexts), 2) # (1,B,2*H)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "        # output (1,B,H)\n",
    "        output = F.log_softmax(self.out(torch.cat((output[0], contexts[0]), 1)),dim=1) # (B,O)\n",
    "        return output, hidden, attn_weights\n",
    "        \n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "one_input = training_pairs[0][0]\n",
    "out = encoder(one_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = AttnDecoderRNN(hidden_size, output_lang.n_words).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "decoder_hidden = decoder.initHidden()\n",
    "encoder_outputs = out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output, hidden, attn_weights = decoder(decoder_input, decoder_hidden, encoder_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.sum(attn_weights,dim=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate=0.01\n",
    "encoder_optimizer = optim.SGD(encoder.parameters(), lr=learning_rate)\n",
    "decoder_optimizer = optim.SGD(decoder.parameters(), lr=learning_rate)\n",
    "criterion = nn.NLLLoss()\n",
    "train_onepair(training_pairs[0][0],training_pairs[0][1],encoder,decoder,encoder_optimizer,decoder_optimizer,criterion, with_attention = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "n_epochs = 50000\n",
    "encoder = EncoderRNN(input_lang.n_words, hidden_size).to(device)\n",
    "decoder = AttnDecoderRNN(hidden_size, output_lang.n_words).to(device)\n",
    "plot_losses,plot_losses_val = trainIters(encoder,decoder,n_epochs,print_every=2500,with_attention = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(running_mean(plot_losses))\n",
    "plt.plot([2500*i for i in range(int(n_epochs/2500))], plot_losses_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_attention(encoder, decoder, sentence, max_length=MAX_LENGTH):\n",
    "    with torch.no_grad():\n",
    "        input_tensor = tensorFromSentence(input_lang, sentence)\n",
    "        input_length = input_tensor.size()[0]\n",
    "        encoder_hidden = encoder.initHidden()\n",
    "\n",
    "        encoder_outputs = encoder(input_tensor)\n",
    "\n",
    "        decoder_input = torch.tensor([[SOS_token]], device=device)  # SOS\n",
    "\n",
    "        decoder_hidden = decoder.initHidden()\n",
    "\n",
    "        decoded_words = []\n",
    "        decoder_attentions = torch.zeros(max_length, max_length)\n",
    "\n",
    "        for di in range(max_length):\n",
    "            decoder_output, decoder_hidden, decoder_attention = decoder(\n",
    "                decoder_input, decoder_hidden, encoder_outputs)\n",
    "            #print(decoder_attention.shape)\n",
    "            decoder_attentions[di,:input_length] = decoder_attention[:,0,0].data\n",
    "            topv, topi = decoder_output.data.topk(1)\n",
    "            if topi.item() == EOS_token:\n",
    "                decoded_words.append('<EOS>')\n",
    "                break\n",
    "            else:\n",
    "                decoded_words.append(output_lang.index2word[topi.item()])\n",
    "\n",
    "            decoder_input = topi.squeeze().detach()\n",
    "\n",
    "        return decoded_words, decoder_attentions[:di + 1,:input_length]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "francais = \"c est un jeune directeur plein de talent .\"#\"elle a cinq ans de moins que moi .\"\n",
    "output_words, attentions = evaluate_attention(encoder, decoder, francais)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attentions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.sum(attentions,dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.matshow(attentions.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexesFromSentence(input_lang, francais)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tensorFromSentence(input_lang, francais)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateRandomly(encoder, decoder, with_attention =True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "def showAttention(input_sentence, output_words, attentions):\n",
    "    # Set up figure with colorbar\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111)\n",
    "    cax = ax.matshow(attentions.numpy(), cmap='bone')\n",
    "    fig.colorbar(cax)\n",
    "\n",
    "    # Set up axes\n",
    "    ax.set_xticklabels([''] + input_sentence.split(' ') +\n",
    "                       ['<EOS>'],rotation =90)\n",
    "    ax.set_yticklabels([''] + output_words)\n",
    "\n",
    "    # Show label at every tick\n",
    "    ax.xaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "    ax.yaxis.set_major_locator(ticker.MultipleLocator(1))\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def evaluateAndShowAttention(input_sentence):\n",
    "    output_words, attentions = evaluate_attention(\n",
    "        encoder, decoder, input_sentence)\n",
    "    print('input =', input_sentence)\n",
    "    print('output =', ' '.join(output_words))\n",
    "    showAttention(input_sentence, output_words, attentions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"elle a cinq ans de moins que moi .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"elle est trop petite .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"je ne crains pas de mourir .\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluateAndShowAttention(\"c est un jeune directeur plein de talent .\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Coming back to the tutorial, the attention mechanism coded there is strange as the attention weights are computed only from the previous prediction $y_{t-1}$ and the hidden state of the decoder $s_{t-1}$. In particular, the attention weights are not computed based on the outputs of the encoder! Indeed the attention weights are always of dimension `MAX_LENGTH` whatever the size of the sentence (and the outputs of the decoder are padded). The illustration below (from the tutorial) illustrates the code:\n",
    "\n",
    "![](https://pytorch.org/tutorials/_images/attention-decoder-network.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AttnDecoderRNN_tuto(nn.Module):\n",
    "    def __init__(self, hidden_size, output_size, dropout_p=0.1, max_length=MAX_LENGTH):\n",
    "        super(AttnDecoderRNN_tuto, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.output_size = output_size\n",
    "        self.dropout_p = dropout_p\n",
    "        self.max_length = max_length\n",
    "\n",
    "        self.embedding = nn.Embedding(self.output_size, self.hidden_size)\n",
    "        self.attn = nn.Linear(self.hidden_size * 2, self.max_length)\n",
    "        self.attn_combine = nn.Linear(self.hidden_size * 2, self.hidden_size)\n",
    "        self.dropout = nn.Dropout(self.dropout_p)\n",
    "        self.gru = nn.GRU(self.hidden_size, self.hidden_size)\n",
    "        self.out = nn.Linear(self.hidden_size, self.output_size)\n",
    "\n",
    "    def forward(self, input, hidden, encoder_outputs):\n",
    "        embedded = self.embedding(input).view(1, 1, -1)\n",
    "        embedded = self.dropout(embedded)\n",
    "\n",
    "        attn_weights = F.softmax(\n",
    "            self.attn(torch.cat((embedded[0], hidden[0]), 1)), dim=1)\n",
    "        \n",
    "        attn_applied = torch.bmm(attn_weights.unsqueeze(0),\n",
    "                                 encoder_outputs.unsqueeze(0))\n",
    "\n",
    "        output = torch.cat((embedded[0], attn_applied[0]), 1)\n",
    "        output = self.attn_combine(output).unsqueeze(0)\n",
    "\n",
    "        output = F.relu(output)\n",
    "        output, hidden = self.gru(output, hidden)\n",
    "\n",
    "        output = F.log_softmax(self.out(output[0]), dim=1)\n",
    "        return output, hidden, attn_weights\n",
    "\n",
    "    def initHidden(self):\n",
    "        return torch.zeros(1, 1, self.hidden_size, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_tuto = AttnDecoderRNN_tuto(hidden_size, output_lang.n_words).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder_input = torch.tensor([[SOS_token]], device=device)\n",
    "decoder_hidden = decoder_tuto.initHidden()\n",
    "encoder_outputs = torch.randn(MAX_LENGTH, hidden_size).to(device)\n",
    "output, hidden, attn_weights = decoder_tuto(decoder_input, decoder_hidden, encoder_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedded = decoder_tuto.embedding(decoder_input)\n",
    "embedded.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden = decoder_hidden\n",
    "hidden.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cat((embedded[0], hidden[0]), 1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_weights = F.softmax(decoder_tuto.attn(torch.cat((embedded[0], hidden[0]), 1)), dim=1)\n",
    "attn_weights.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_weights.unsqueeze(0).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "encoder_outputs.unsqueeze(0).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_applied = torch.bmm(attn_weights.unsqueeze(0),\n",
    "                                 encoder_outputs.unsqueeze(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "attn_applied.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Dataflowr](https://raw.githubusercontent.com/dataflowr/website/master/_assets/dataflowr_logo.png)](https://dataflowr.github.io/website/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dldiy",
   "language": "python",
   "name": "dldiy"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
